#include <stdio.h>
#include <math.h>
#include <stdlib.h>

#include "crypto/types.h"
#include "crypto/rsa.h"
#include "crypto/tool.h"

static void ex_gcd(uint64 e, int m, uint64 *d, uint64 *y)
{
    if (m == 0)
    {
        *d = 1;
        *y = 0;
        return;
    }
    ex_gcd(m, e % m, y, d);
    (*y) -= e / m * (*d);
}

// void int ex_gcd(int a, int b, int *x, int *y)
// {
//     if (b == 0)
//     {
//         x = 1;
//         y = 0;
//         return a;
//     }
//     int d = ex_gcd(b, a % b, y, x);
//     (*y) -= (a / b) * (*x);
//     return d;
// }

static int gcd(int a, int b)
{
    if (b)
    {
        return gcd(b, a % b);
    }
    else
    {
        return a;
    }
}

static int lcm(int a, int b)
{
    return a / gcd(a, b) * b; // 防溢出 ， 很妙啊 ，大家可以记一下
}

static bool prime(int p)
{
    if (p <= 17)
    {
        return 0;
    }
    if (p % 6 != 1 && p % 6 != 5)
    {
        return 0;
    }
    int s_half = sqrt(p);
    for (int i = 5; i <= s_half; i += 6)
    {
        if (p % i == 0 || p % (i + 2) == 0)
        {
            return 0;
        }
    }
    return 1;
}

uint powMod(uint base, uint64 e, uint mod)
{
    uint res = 1 % mod;
    uint64 t = base;
    while (e)
    {
        if (e & 1) // 此处等价于if(n%2==1)  if(n&0)等价于n%2==0
            res = res * t % mod;
        t = t * t % mod;
        e >>= 1; // 等价于n=n/2
    }
    return res;
}

static bool judgeDEM(uint64 d, uint e, uint m)
{
    return (d * e) % m == 1;
}

RSASuit *RSAgetRSASuit(RandomFunc randomFunc)
{
    RSASuit *suit = (RSASuit *)malloc(sizeof(RSASuit));
    uint e, m;
    uint64 y, d;
    uint16 tmp16;
    suit->e = RSA_DEFAULT_E;
    e = RSA_DEFAULT_E;
    do
    {
        randomBuf(&tmp16, 2, randomFunc);
        if (tmp16 & 0x1111111100000000 == 0)
        {
            tmp16 = tmp16 | (~tmp16);
        }
        tmp16 = tmp16 / 6 * 6 + 5;
        while (!prime(tmp16))
        {
            tmp16 += 6;
        }
        suit->p = tmp16;
        randomBuf(&tmp16, 2, randomFunc);
        if (tmp16 & 0x1111111100000000 == 0)
        {
            tmp16 = tmp16 | (~tmp16);
        }
        tmp16 = tmp16 / 6 * 6 + 5;
        while (!prime(tmp16))
        {
            tmp16 += 6;
        }
        suit->q = tmp16;
        suit->n = suit->p * suit->q;
        // suit->m = lcm(suit->p - 1, suit->q - 1);
        suit->m = (suit->p - 1) * (suit->q - 1);
        m = suit->m;
        ex_gcd(e, m, &d, &y);
        suit->d = d;
    } while (suit->m % RSA_DEFAULT_E == 0 || suit->d >= RSA_MAX_D || !judgeDEM(suit->d, suit->e, suit->m));

    return suit;
}

SAFE_array *RSAEncode(uint N, uint64 E, SAFE_array *plainText)
{
    // printf("\nRSAEncode : N: %d E: %lld\n", N, E);
    int len = plainText->size;
    SAFE_array *sa = (SAFE_array *)malloc(offset(SAFE_array, data) + len * 4);
    for (int i = 0; i < len; i++)
    {
        *(int *)((sa->data) + i * 4) = powMod(*(uint8 *)(plainText->data + i), E, N);
    }
    len *= 4;
    sa->size = len;
    sa->cap = len;
    return sa;
}

SAFE_array *RSADecode(uint N, uint64 D, SAFE_array *cryptoText)
{
    // printf("\nRSADecode : N: %d D: %lld\n", N, D);
    int len = cryptoText->size / 4;
    SAFE_array *sa = (SAFE_array *)malloc(offset(SAFE_array, data) + len);
    for (int i = 0; i < len; i++)
    {
        uint ret = (uint)powMod(*(uint *)(cryptoText->data + i * 4), D, N);
        *((uint8 *)(sa->data + i)) = (uint8)ret;
    }
    sa->size = len;
    sa->cap = len;
    return sa;
}

void RSASuitPrint(RSASuit *suit)
{
    printf("*************RSA key**************\n");
    printf("public key:\t%u\n", suit->e);
    printf("private key:\t%llu\n", suit->d);
    printf("N:\t%u\n", suit->n);
    printf("m:\t%u\n", suit->m);
    printf("p:\t%u\n", suit->p);
    printf("q:\t%u\n", suit->q);
    printf("**********************************\n");
}